import numpy as np
import os
import pickle


def read_pickle(file_path, suffix='.pkl'):  # 读取 .pkl 文件
    """
    功能：从.pkl文件(file_path)中读取数据，并向返回值data里写入数据的函数
    """
    assert os.path.splitext(file_path)[1] == suffix
    with open(file_path, 'rb') as f:
        data = pickle.load(f)  # pickle库的load()函数来反序列化.pkl文件，把二进制数据恢复成Python对象。
    return data


def write_pickle(results, file_path):
    """
    功能：将对象results写入到.pkl文件中(file_path)
    """
    with open(file_path, 'wb') as f:
        pickle.dump(results, f)  # pickle.dump()函数将results这个Python对象进行序列化，并写入到f指向的文件中。


def read_points(file_path, dim=4):
    """
    功能：读取file_path指向的点云.bin文件，并转换为np数组，形状为(N,4)
    """
    suffix = os.path.splitext(file_path)[1]  # 使用os.path.splitext()函数分割文件路径，[1]提取文件的扩展名部分
    assert suffix in ['.bin', '.ply']
    if suffix == '.bin':
        return np.fromfile(file_path, dtype=np.float32).reshape(-1, dim)  # 如果后缀名是.bin，则读取点云文件，并转换为np数组，形状为(N,4)
    else:
        raise NotImplementedError


def write_points(lidar_points, file_path):
    """
    将lidar_points的数据保存到file_path指向的文件中
    """
    suffix = os.path.splitext(file_path)[1]
    assert suffix in ['.bin', '.ply']
    if suffix == '.bin':
        with open(file_path, 'wb') as f:  # 存储lidar_points的数据到file_path内
            lidar_points.tofile(f)
    else:
        raise NotImplementedError


def read_calib(file_path, extend_matrix=True):
    """
    功能：从 KITTI数据集中的校准文件（calib.txt）中读取相机、激光雷达、IMU等传感器的内外参矩阵，并返回一个字典 calib_dict，供后续3D点云转换到相机坐标系使用。
    参数：calib.txt 文件路径。
    过程：读取文件每行 → 提取参数 → 转换为 numpy 数组 → 扩展为4×4矩阵。
    返回值：一个字典 calib_dict，包含了相机、激光雷达、IMU的内外参矩阵。
    """
    with open(file_path, 'r') as f:
        lines = f.readlines()
    lines = [line.strip() for line in lines]  # 按行读取lines，同时去除每一行的前后空格和换行，并填入列表lines

    # 读取P1、P2、P3、P4相机的投影矩阵
    P0 = np.array([item for item in lines[0].split(' ')[1:]], dtype=np.float32).reshape(3, 4)
    P1 = np.array([item for item in lines[1].split(' ')[1:]], dtype=np.float32).reshape(3, 4)
    P2 = np.array([item for item in lines[2].split(' ')[1:]], dtype=np.float32).reshape(3, 4)
    P3 = np.array([item for item in lines[3].split(' ')[1:]], dtype=np.float32).reshape(3, 4)

    # R0_rect：校正矩阵，对多传感器数据进行旋转修正，形状为 3×3。
    R0_rect = np.array([item for item in lines[4].split(' ')[1:]], dtype=np.float32).reshape(3, 3)
    # Tr_velo_to_cam：激光雷达到相机的变换矩阵，形状为 3×4。
    Tr_velo_to_cam = np.array([item for item in lines[5].split(' ')[1:]], dtype=np.float32).reshape(3, 4)
    # Tr_imu_to_velo：IMU到激光雷达的变换矩阵，形状为 3×4。
    Tr_imu_to_velo = np.array([item for item in lines[6].split(' ')[1:]], dtype=np.float32).reshape(3, 4)

    if extend_matrix:  # 判断是否需要将这些矩阵扩展为4×4齐次变换矩阵（默认是 True）。
        P0 = np.concatenate([P0, np.array([[0, 0, 0, 1]])], axis=0)
        P1 = np.concatenate([P1, np.array([[0, 0, 0, 1]])], axis=0)
        P2 = np.concatenate([P2, np.array([[0, 0, 0, 1]])], axis=0)
        P3 = np.concatenate([P3, np.array([[0, 0, 0, 1]])], axis=0)

        R0_rect_extend = np.eye(4, dtype=R0_rect.dtype)  # 创建单位矩阵
        R0_rect_extend[:3, :3] = R0_rect  # 拓展矩阵的左上角被赋值为R0_rect
        R0_rect = R0_rect_extend  # 引用赋值，使R0_rect和R0_rect_extend指向的是同一个(4,4)的矩阵

        Tr_velo_to_cam = np.concatenate([Tr_velo_to_cam, np.array([[0, 0, 0, 1]])], axis=0)
        Tr_imu_to_velo = np.concatenate([Tr_imu_to_velo, np.array([[0, 0, 0, 1]])], axis=0)

    calib_dict = dict(
        P0=P0,
        P1=P1,
        P2=P2,
        P3=P3,
        R0_rect=R0_rect,
        Tr_velo_to_cam=Tr_velo_to_cam,
        Tr_imu_to_velo=Tr_imu_to_velo
    )
    return calib_dict


def read_label(file_path):
    """
    功能：从指定的 KITTI label 文件中读取 3D 目标检测的标注信息，并将其转换为字典形式返回，便于后续处理和使用。
    输入:的是一个 label 文件路径。
    输出:是一个包含所有标注信息的字典。
    """
    with open(file_path, 'r') as f:
        lines = f.readlines()
    lines = [line.strip().split(' ') for line in lines]  # 去除前后空格和换行符，生成二维列表（行和列可以不规则）
    annotation = {}
    annotation['name'] = np.array([line[0] for line in lines])  # 所有真值对象的类别,shape(N,)
    annotation['truncated'] = np.array([line[1] for line in lines], dtype=np.float32)  # 截断比例(N,)
    annotation['occluded'] = np.array([line[2] for line in lines], dtype=np.int32)  # 遮挡等级(N,)
    annotation['alpha'] = np.array([line[3] for line in lines], dtype=np.float32)  # 解析观测角度(N,)
    annotation['bbox'] = np.array([line[4:8] for line in lines], dtype=np.float32)  # 解析2D边界框(N,4)
    annotation['dimensions'] = np.array([line[8:11] for line in lines], dtype=np.float32)[:,
                               [2, 0, 1]]  # 3D尺寸hwl -> 相机坐标系 (lhw)(N,3)
    annotation['location'] = np.array([line[11:14] for line in lines], dtype=np.float32)  # 解析3D位置，物体在相机坐标系下的3D坐标位置(N,3)
    annotation['rotation_y'] = np.array([line[14] for line in lines], dtype=np.float32)  # 物体绕 y 轴的旋转角度(N,)

    return annotation


def write_label(result, file_path, suffix='.txt'):
    '''
    result: dict,
    file_path: str
    '''
    assert os.path.splitext(file_path)[1] == suffix
    name, truncated, occluded, alpha, bbox, dimensions, location, rotation_y, score = \
        result['name'], result['truncated'], result['occluded'], result['alpha'], \
            result['bbox'], result['dimensions'], result['location'], result['rotation_y'], \
            result['score']

    with open(file_path, 'w') as f:
        for i in range(len(name)):
            bbox_str = ' '.join(map(str, bbox[i]))
            hwl = ' '.join(map(str, dimensions[i]))
            xyz = ' '.join(map(str, location[i]))
            line = f'{name[i]} {truncated[i]} {occluded[i]} {alpha[i]} {bbox_str} {hwl} {xyz} {rotation_y[i]} {score[i]}\n'
            f.writelines(line)
